
#include <png.h>
#include <BImage>
#include <BFile>
#include "ReadWrite_png.h"

ReadWrite_png::ReadWrite_png()
{
	_author = "zym";
	_description = "This is BWE Plugin for reading *.png image file, write width libpng.";
	_version = "1.0";
}
ReadWrite_png::~ReadWrite_png()
{

}

void ReadWrite_png::setOptions(const BString& options)
{
	_options = options;
}
const BString& ReadWrite_png::options() const
{
	return _options;
}

const BString& ReadWrite_png::author() const
{
	return _author;
}
const BString& ReadWrite_png::version() const
{
	return _version;
}
const BString& ReadWrite_png::description() const
{
	return _description;
}
const BString& ReadWrite_png::message() const
{
	return _message;
}

BObject* ReadWrite_png::load(const BString& fileName)
{
	BString name = fileName.name();
	name.remove(name.find('.'), name.size());
	BImage* image = new BImage(name);
	if (!load(fileName, image))
	{
		delete image;
		return 0;
	}
	return image;
}
bool ReadWrite_png::load(const BString& fileName, BObject* object)
{
	_message.clear();
	BImage* image = dynamic_cast<BImage*>(object);
	if (!image)
	{
		_message << "Read object is not image but " << object->className();
		return false;
	}

	FILE* fp = fopen(fileName.cstr(), "rb");
	if (!fp)
	{
		_message << "can't open file " << fileName;
		return false;
	}

	png_byte header[8];
	fread(header, 1, 8, fp);
	bool is_png = !png_sig_cmp(header, 0, 8);
	if(!is_png)
	{
		fclose(fp);
		_message << "file content is not valid for png : " << fileName;
		return false;
	}

	png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
    if(!png_ptr)
	{
		fclose(fp);
		return false;
	}

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr)
    {
        png_destroy_read_struct(&png_ptr, 0, 0);
		fclose(fp);
		return false;
    }

	png_init_io(png_ptr, fp);
	png_set_sig_bytes(png_ptr, 8);
	png_read_info(png_ptr, info_ptr);

	int width            = png_get_image_width(png_ptr, info_ptr);
    int height           = png_get_image_height(png_ptr, info_ptr);
    int bit_depth        = png_get_bit_depth(png_ptr, info_ptr);
    int color_type       = png_get_color_type(png_ptr, info_ptr);
    int filter_method    = png_get_filter_type(png_ptr, info_ptr);
    int compression_type = png_get_compression_type(png_ptr, info_ptr);
    int interlace_type   = png_get_interlace_type(png_ptr, info_ptr);
	int channels = png_get_channels(png_ptr, info_ptr);
	int rowbytes = (int)png_get_rowbytes(png_ptr, info_ptr);

	if (bit_depth == 16)
		png_set_strip_16(png_ptr);
	if (bit_depth < 8)
		png_set_packing(png_ptr);

	Format format = Format_None;
	switch (color_type)
	{
	case PNG_COLOR_TYPE_PALETTE:
		png_set_palette_to_rgb(png_ptr);
		if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
		{
			png_set_tRNS_to_alpha(png_ptr);
			format = Format_RGBA;
		}
		else
		{
			format = Format_RGB;
		}
		break;
	case PNG_COLOR_TYPE_RGB:
		format = Format_RGB;
		break;
	case PNG_COLOR_TYPE_RGBA:
		format = Format_RGBA;
		break;
	case PNG_COLOR_TYPE_GRAY:
		format = Format_Alpha;
		if (bit_depth < 8)
			png_set_expand_gray_1_2_4_to_8(png_ptr);
		break;
	case PNG_COLOR_TYPE_GRAY_ALPHA:
		format = Format_Gray_Alpha;
		break;
	default:
		png_destroy_read_struct(&png_ptr, &info_ptr, 0);
		fclose(fp);
		_message << "invalid format : " << color_type;
		return false;
	}

	image->setFormat(format);
	image->setWidth(width);
	image->setHeight(height);

	BArray<BByte*> pointers(height);
	for (int row = 0; row < height; row++)
	{
		pointers[row] = image->bytes(row);
	}

	png_read_image(png_ptr, pointers.data());
	png_destroy_read_struct(&png_ptr, &info_ptr, 0);
	fclose(fp);

	return true;
}
bool ReadWrite_png::save(const BString& fileName, const BObject* object) const
{
	_message.clear();
	const BImage* image = dynamic_cast<const BImage*>(object);
	if (!image)
	{
		_message << " Save object is not image but " << object->className();
		return false;
	}

	FILE* fp = fopen(fileName.cstr(), "wb");
	if (!fp)
	{
		_message << "can't open file " << fileName;
		return false;
	}

	png_struct* png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (!png_ptr)
	{
		_message << "create write struct failed.";
		return false;
	}

	png_info* info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		_message << "create info struct failed.";
		png_destroy_write_struct(&png_ptr, NULL);
		return false;
	}

	png_init_io(png_ptr, fp);

	Format format = image->format();
	int width = image->width();
	int height = image->height();
	const BByte* data = image->pixels()->data();

	BArray<png_bytep> row_pointers;
	switch (format)
	{
	case Format_RGBA:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 4);
			row_pointers[r] = bytes;
		}
		break;
	case Format_RGB:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 3);
			row_pointers[r] = bytes;
		}
		break;
	case Format_Gray_Alpha:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_GRAY_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 2);
			row_pointers[r] = bytes;
		}
		break;
	case Format_Gray:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_GRAY, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 1);
			row_pointers[r] = bytes;
		}
		break;
	default:
		png_destroy_write_struct(&png_ptr, &info_ptr);
		fclose(fp);
		return false;
	}

	png_set_packing(png_ptr);
	png_write_info(png_ptr, info_ptr);
	png_write_image(png_ptr, row_pointers.data());
	png_write_end(png_ptr, NULL);
	png_destroy_write_struct(&png_ptr, &info_ptr);
	fclose(fp);
	return true;
}

static void pngReadCallback(png_structp png_ptr, png_bytep data, png_size_t length)
{
	BBuffer* buffer = (BBuffer*)png_get_io_ptr(png_ptr);
	buffer->read(data, (int)length);
}
static void pngWriteCallback(png_structp png_ptr, png_bytep data, png_size_t length)
{
	BBuffer* buffer = (BBuffer*)png_get_io_ptr(png_ptr);
	buffer->write(data, (int)length);
}

BObject* ReadWrite_png::read(const BBuffer& buffer)
{
	BImage* image = new BImage();
	if (!read(buffer, image))
	{
		delete image;
		return 0;
	}
	return image;
}
bool ReadWrite_png::read(const BBuffer& buffer, BObject* object)
{
	BImage* image = dynamic_cast<BImage*>(object);
	if (!image)
	{
		_message << "Read object is not image but " << object->className();
		return false;
	}

	png_byte header[8];
	buffer.read(header, 8);
	bool is_png = !png_sig_cmp(header, 0, 8);
	if (!is_png)
	{
		_message << "Read image failed: bytes is not png.";
		return false;
	}
	buffer.seek(0);

	png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, 0, 0, 0);
	if (!png_ptr)
	{
		_message << "Read image failed: create read struct failed.";
		return false;
	}

	png_infop info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		png_destroy_read_struct(&png_ptr, 0, 0);
		_message << "Read image failed: create info struct failed.";
		return false;
	}

	png_set_read_fn(png_ptr, (png_voidp)&buffer, pngReadCallback);

	//png_set_sig_bytes(png_ptr, 8);
	png_read_info(png_ptr, info_ptr);

	int width = png_get_image_width(png_ptr, info_ptr);
	int height = png_get_image_height(png_ptr, info_ptr);
	int bit_depth = png_get_bit_depth(png_ptr, info_ptr);
	int color_type = png_get_color_type(png_ptr, info_ptr);
	int filter_method = png_get_filter_type(png_ptr, info_ptr);
	int compression_type = png_get_compression_type(png_ptr, info_ptr);
	int interlace_type = png_get_interlace_type(png_ptr, info_ptr);
	int channels = png_get_channels(png_ptr, info_ptr);
	int rowbytes = (int)png_get_rowbytes(png_ptr, info_ptr);

	if (bit_depth == 16)
		png_set_strip_16(png_ptr);

	Format format = Format_None;
	switch (color_type)
	{
	case PNG_COLOR_TYPE_PALETTE:
		png_set_packing(png_ptr);
		png_set_palette_to_rgb(png_ptr); //Expand data to 24-bit RGB or 32-bit RGBA if alpha available.
		if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
		{
			png_set_tRNS_to_alpha(png_ptr);
			format = Format_RGBA;
			rowbytes *= 4;
		}
		else
		{
			format = Format_RGB;
			rowbytes *= 3;
		}
		break;
	case PNG_COLOR_TYPE_RGB:
		format = Format_RGB;
		break;
	case PNG_COLOR_TYPE_RGBA:
		format = Format_RGBA;
		break;
	case PNG_COLOR_TYPE_GRAY:
		format = Format_Alpha;
		if (bit_depth < 8)
			png_set_expand_gray_1_2_4_to_8(png_ptr);
		break;
	case PNG_COLOR_TYPE_GRAY_ALPHA:
		format = Format_Gray_Alpha;
		break;
	default:
		_message << "Read image failed: png format is not supported.";
		png_destroy_read_struct(&png_ptr, &info_ptr, 0);
		return false;
	}

	BHolder<BByteArray> pixels = new BByteArray(rowbytes * height);
	BByte* data = pixels->data();
	BArray<BByte*> row_pointers(height);
	for (int row = 0; row < height; row++)
	{
		row_pointers[row] = &data[rowbytes * row];
	}

	png_read_image(png_ptr, row_pointers.data());
	png_destroy_read_struct(&png_ptr, &info_ptr, 0);

	image->setFormat(format);
	image->setWidth(width);
	image->setHeight(height);
	image->setPixels(pixels.ptr());

	return true;
}
bool ReadWrite_png::write(BBuffer& buffer, const BObject* object) const
{
	const BImage* image = dynamic_cast<const BImage*>(object);
	if (!image)
	{
		_message << "\tRead object is not image but " << object->className();
		return false;
	}

	const BByteArray* pixels = image->pixels();
	if (!pixels)
		return false;

	png_struct* png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if (!png_ptr)
		return false;

	png_info* info_ptr = png_create_info_struct(png_ptr);
	if (!info_ptr)
	{
		png_destroy_write_struct(&png_ptr, NULL);
		return false;
	}

	png_set_write_fn(png_ptr, &buffer, pngWriteCallback, 0);

	Format format = image->format();
	int width = image->width();
	int height = image->height();
	const BByte* data = pixels->data();

	BArray<png_bytep> row_pointers;
	switch (format)
	{
	case Format_RGBA:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 4);
			row_pointers[r] = bytes;
		}
		break;
	case Format_RGB:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 3);
			row_pointers[r] = bytes;
		}
		break;
	case Format_Gray_Alpha:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_GRAY_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 2);
			row_pointers[r] = bytes;
		}
		break;
	case Format_Gray:
		png_set_IHDR(png_ptr, info_ptr, width, height, 8, PNG_COLOR_TYPE_GRAY, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);
		row_pointers.resize(height);
		for (int r = 0; r < height; r++)
		{
			BByte* bytes = (BByte*)data + (r * width * 1);
			row_pointers[r] = bytes;
		}
		break;
	default:
		png_destroy_write_struct(&png_ptr, &info_ptr);
		return false;
	}

	png_set_packing(png_ptr);
	png_write_info(png_ptr, info_ptr);
	png_write_image(png_ptr, row_pointers.data());
	png_write_end(png_ptr, NULL);
	png_destroy_write_struct(&png_ptr, &info_ptr);

	return true;
}
